'''
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Created on: 7/26/2009

Version: 0.1

@author: bryan

Utilities to support basic JointTrack functions.
'''
import JT

import csv
import os
import numpy as np

def mag(vector):
    '''
    return the magnitude of the given vector
    '''
    if type(vector).__name__ in ('int','float'):
        return vector
    m = 0
    for i in vector:
        m += i**2
    return m**0.5

def genericImageReader(filename):
    """
    A generic VTK image reader.
    filename: a path 
        the appropriate reader will be selected based on the filename extension.
    
    output: VTK image data (reader.GetOutput())
    """
    ext = filename[-3:].lower()
    if (ext == "mha" or ext == 'mhd'):
        # mha image can handle float pixel values
        reader = JT.vtk.vtkMetaImageReader()
    elif (ext == "png"):
        # png image can can be UC or US
        reader = JT.vtk.vtkPNGReader()
    elif (ext == "bmp"):
        # bmp image will be written as UC regardless of input type
        reader = JT.vtk.vtkBMPReader()
    elif (ext == "jpg"):
        # jpg image must be UC
        reader = JT.vtk.vtkJPEGReader()
    elif (ext == "tiff" or ext == "tif"):
        # jpg image must be UC
        reader = JT.vtk.vtkTIFFReader()
    else:
         print "ERROR: No reader implemented for %s files, nothing done" % ext

    reader.SetFileName(filename)
    try:
        reader.Update()
        return reader.GetOutput()
    except:
        print "Error reading file: %s" % filename

def genericImageWriter(image, filename):
    """
    A generic VTK image writer.
    image: VTK image data
    filename: a path 
        the appropriate writer will be selected based on the filename extension.
    """
    ext = filename[-3:].lower()
    if (ext == "mha" or ext == 'mhd'):
        # mha image can handle float pixel values
        writer = JT.vtk.vtkMetaImageWriter()
        writer.SetCompression(0)
    elif (ext == "png"):
        # png image can can be UC or US
        writer = JT.vtk.vtkPNGWriter()
    elif (ext == "bmp"):
        # bmp image will be written as UC regardless of input type
        writer = JT.vtk.vtkBMPWriter()
    elif (ext == "jpg"):
        # jpg image must be UC
        writer = JT.vtk.vtkJPEGWriter()
    elif (ext == "tiff" or ext == "tif"):
        # jpg image must be UC
        writer = JT.vtk.vtkTIFFWriter()
    else:
         print "ERROR: No writer implemented for %s files, nothing done" % ext

    writer.SetFileName(filename)
    writer.SetInput(image)
    try:
        writer.Write()
    except:
        print "Error writing file: %s" % filename
        
def vtk2numpy(vtk_image):
    """
    A convenience function for using numpy_support to convert a vtk image to
    a numpy array. 
    """
    numpy_image = JT.numpy_support.vtk_image_to_numpy(vtk_image)
    return numpy_image


def GetTransform(pose):
        """
        Given a 6 element vector (Tx,Ty,Tx,Rx,Ry,Rz) return a transformation
        matrix.  (NOT shifted or concatenated with anything else)
        """
        Tx,Ty,Tz,Rx,Ry,Rz = pose
        # Tcs is TStartTransformation (body center to start pose)
        Tcs = np.eye(4)
        Tcs[0,0] = np.cos(Ry)*np.cos(Rz)
        Tcs[0,1] = -np.cos(Rx)*np.sin(Rz) + np.cos(Rz)*np.sin(Rx)*np.sin(Ry)
        Tcs[0,2] = np.sin(Rx)*np.sin(Rz) + np.cos(Rx)*np.cos(Rz)*np.sin(Ry)
        Tcs[0,3] = Tx
        Tcs[1,0] = np.cos(Ry)*np.sin(Rz)
        Tcs[1,1] = np.cos(Rx)*np.cos(Rz) + np.sin(Rx)*np.sin(Ry)*np.sin(Rz)
        Tcs[1,2] = -np.cos(Rz)*np.sin(Rx) + np.cos(Rx)*np.sin(Ry)*np.sin(Rz)
        Tcs[1,3] = Ty
        Tcs[2,0] = -np.sin(Ry)
        Tcs[2,1] = np.cos(Ry)*np.sin(Rx)
        Tcs[2,2] = np.cos(Rx)*np.cos(Ry)
        Tcs[2,3] = Tz
        
        return Tcs
    
def GetVTKMatrixFromNumpyArray(npArray):
        """
        Return a VTK matrix from a numpy 4x4 array.
        """
        if (npArray.shape != (4,4)):
            print "Error: numpy array is not 4x4, cannot convert to VTK matrix"
            return
        
        vtkMatrix = JT.vtk.vtkMatrix4x4()
        
        for row in range(npArray.shape[0]):
            for col in range(npArray.shape[1]):
                vtkMatrix.SetElement(row,col,npArray[row][col])
        
        return vtkMatrix
    
def GetNumpyArrayFromVTKMatrix(vtkMatrix):
    """
    Return a numpy array from a VTK Matrix
    """
    try:
        # Make sure that the user suppies a valid VTK matrix
        # This might also work if a transform is given
        vtkMatrix.GetClassName() == 'vtkMatrix4x4'
    except:
        print "Input is not recognized as a 'vtkMatrix4x4'"
        return
    
    npArray = np.eye(4)
    
    for row in range(npArray.shape[0]):
        for col in range(npArray.shape[1]):
            npArray[row][col] = vtkMatrix.GetElement(row,col,)
    
    return npArray

def FindMTre(pose,proj_norm=None,ground_truth=[0,0,0,0,0,0]):
    """
    Given a pose and the ground truth pose, both in 1D vector format 
    (Tx,Ty,Tz,Rx,Ry,Rz) with rotations in radians, this function will 
    return the mTRE as a float.
    
    The mTRE is one way to measure the accuracy of a 
    registration algorithm.  For the Gold Standard data, the ground truth pose
    is [0,0,0,0,0,0].  
    
    mTRE is calculated by sampling a grid of points from within a clinically 
    relevant region.  This is specified as a region 95x45x95 pixels 
    (0.87mm^3/voxel).  This grid is sampled every 10 points to speed up the 
    algorithm (500 points vs 400,000).
    
    If the mTRE in the projection direction (mTREproj) is desired, the optional
    parameter projection_normal should be supplied.  The projection_normal is 
    a unit vector that describes the direction from the detector to the source. 
    
    """
    Tgold_wcs = GetTransform(ground_truth)
    Tgold_wcs_inv = np.linalg.inv(Tgold_wcs)

    Treg_vol = GetTransform(pose)
    Treg_wcs = np.dot(Treg_vol,Tgold_wcs)
    Treg_wcs_inv = np.linalg.inv(Treg_wcs)
    
    # n is the unit vector projecting out of the imaging plane toward
    # the source
    if proj_norm != None:
        n = proj_norm
        
    err = 0
    p_err = 0
    i = 0
    for x in range(-47,48,10):
        for y in range(-22,23,10):
            for z in range(-47,48,10):
                p = np.array([x,y,z,1]) * 0.87
                p[3] = 1
                Preg = np.dot(Treg_wcs_inv,p)
                Pgold = np.dot(Tgold_wcs_inv,p)
                err = err + np.sqrt(((Preg - Pgold)**2).sum())
                if proj_norm != None:
                    p_err = p_err + np.sqrt((np.dot(Preg[0:3] - Pgold[0:3],n)**2).sum())
                i = i + 1
    
    mTRE = err/i
    mTRE_proj = p_err/i
    print "mTRE = %f" % (mTRE)
    
    if proj_norm is not None:
        # if a projection normal is given, return the mTRE in projection 
        # direction, otherwise only give back the mTRE
        return mTRE, mTRE_proj
    else:
        return mTRE



class postProcess(object):
    """
    a class to do some post processing on result data from registrations.
    Currently, it will read in a list of poses (Tx,Ty,Tz,Rx,Ry,Rz) and 
    calculate the mTRE (mean Target Registration Error).  Assumes that the 
    ground truth pose is at [0,0,0,0,0,0].
    """
    def __init__(self,filepath):
        [path_dir, filename] = os.path.split(filepath) 
        [basename, ext] = os.path.splitext(filename)
        open_file = open(filepath,'r')
        self.reader = csv.reader(open_file,delimiter='\t')
        save_path = os.path.join(path_dir,basename + '_out.csv')
        self.save_file = open(save_path,'w')
        self.writer = csv.writer(self.save_file)
    
    def read_data(self):
        """
        """
        data = []
        for line in self.reader:
            for element in line:
                try:
                    data.append(float(element))
                except:
                    data.append(9999)
            
        self.data = np.asarray(data,'f')
        self.data = self.data.reshape(-1,6)
        print self.data.shape
        
        
    def mTRE(self):
        """
        """
        mTRE = np.zeros(len(self.data))
        for pose in self.data:
            # do calculation and write to file in one step
            [mTRE,mTRE_proj] = FindMTre(pose,proj_normal=[])
            self.writer.writerow([])
        
        self.save_file.close()
        
    
    def write_data(self):
        """
        done in mTRE
        """
        
        pass
        
    
    def run(self):
        self.read_data()
        self.mTRE()
